# coding=utf-8
# Chern-C
# Copyright (C) 2024 NJU-SAML
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import argparse
import json
import os
from typing import Any, Dict, List, Optional, Tuple
from pathlib import Path
import yaml
from chernc.config import CherncConfig
from chernc.constants import DEFAULT_LOGGER_DIR, DEFAULT_LOGGER_FILE_NAME
from chernc.logging import build_logger
from chernc.agents.translator import PoolBasedModelTranslator
from chernc.core.transpiler import ProjectTranspiler
from chernc.utils.hash_utils import calculate_md5
from chernc.llm.api_inference import AsyncAPIInference

logger = build_logger(__name__ + ".c2rust", DEFAULT_LOGGER_FILE_NAME, logger_dir=DEFAULT_LOGGER_DIR)

logger = build_logger("chernc.c2rust", DEFAULT_LOGGER_FILE_NAME, logger_dir=DEFAULT_LOGGER_DIR)


def main(
    config_path: str,
    input_dir: str,
    output_dir: str,
    codeql: str = "codeql",
    codeql_db: str = "./database",
):
    # Read Configs
    try:
        chernc_config_dict: Dict[str, Any] = yaml.safe_load(Path(config_path).read_text())
    except FileNotFoundError as error:
        message = "Error: yml config file not found."
        logger.exception(message)
        raise FileNotFoundError(error, message) from error

    logger.info("Chern-C configs: \n" + json.dumps(chernc_config_dict, indent=2))

    chernc_config: CherncConfig = CherncConfig.model_validate(chernc_config_dict)

    # Setup LLM Endpoints
    pool_translator = PoolBasedModelTranslator()
    endpoints = chernc_config.endpoints
    for llm_endpoint in endpoints:
        for apikey in llm_endpoint.apikeys:
            pool_translator.register_a_model(
                llm_endpoint.model + ":" + apikey,
                AsyncAPIInference(
                    llm_endpoint.model,
                    api_key=apikey,
                    base_url=llm_endpoint.baseurl,
                    rate_limit=llm_endpoint.rate_limit,
                ),
                llm_endpoint.features,
            )

    # Setup Transpiler
    project_hash = calculate_md5(os.path.abspath(input_dir))
    os.makedirs(codeql_db, exist_ok=True)
    database_path = os.path.join(codeql_db, project_hash)

    project = ProjectTranspiler(
        config=chernc_config,
        project_path=input_dir,
        database_path=database_path,
        translator=pool_translator,
        codeql_bin=codeql,
    )

    # Translate the project
    project.translate_project(output_dir)


def cli():
    parser = argparse.ArgumentParser(description="ChernC2Rust - Translate C/CPP project to Rust language.")
    parser.add_argument("-c", "--config", default="./chernc_configs.yml", help="The config of ChernC2rust")
    parser.add_argument("-i", "--input", required=True, help="Input directory containing C/CPP results")
    parser.add_argument("-o", "--output", required=True, help="Output directory to save C/CPP results")
    parser.add_argument("--codeql", default="codeql", help="The codeql execution binary path")
    parser.add_argument("--codeql-db", default="./database", help="The codeql cache dataset path")

    args = parser.parse_args()

    # 确保输入和输出目录存在
    if not os.path.isdir(args.input):
        raise ValueError(f"Input directory '{args.input}' does not exist.")
    if not os.path.isdir(args.output):
        os.makedirs(args.output, exist_ok=True)
        # raise ValueError(f"Output directory '{args.output}' does not exist.")

    main(args.config, args.input, args.output, args.codeql, args.codeql_db)


if __name__ == "__main__":
    cli()
